{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f866cc0a-460e-4dbf-91b7-f541a4c0eda6",
   "metadata": {
    "tags": []
   },
   "source": [
    "# FLARE Query Engine\n",
    "\n",
    "Adapted from the paper \"Active Retrieval Augmented Generation\"\n",
    "\n",
    "Currently implements FLARE Instruct, which tells the LLM to generate retrieval instructions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "920e6742-5296-46ed-83f5-8b0d999b88a3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index.query_engine import FLAREInstructQueryEngine\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    StorageContext,\n",
    "    LLMPredictor,\n",
    "    ServiceContext,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "507256f4-7831-4ad3-8bdd-c19ad6623d6d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "service_context = ServiceContext.from_defaults(\n",
    "    llm=OpenAI(model=\"gpt-4\", temperature=0), chunk_size=512\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5dac4df9-a42f-4cc5-a708-542e3c33f3d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "documents = SimpleDirectoryReader(\"../data/paul_graham\").load_data()\n",
    "index = VectorStoreIndex.from_documents(documents, service_context=service_context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a749d71-2ad5-4ee3-aa1e-48e9567b706d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "index_query_engine = index.as_query_engine(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "91fcdbcd-f0f6-4580-9016-d89394e05aab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "flare_query_engine = FLAREInstructQueryEngine(\n",
    "    query_engine=index_query_engine,\n",
    "    service_context=service_context,\n",
    "    max_iterations=7,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6a6dcb43-8854-481a-8d6e-62e2c355b72c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<llama_index.query_engine.flare.base.FLAREInstructQueryEngine at 0x1475ba7a0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "flare_query_engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e1ca238e-884b-4221-b3a8-0bb106fa0de9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mQuery: Can you tell me about the author's trajectory in the startup world?\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: The author began their journey in the startup world by [Search(What did the author do in the startup world?)]\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: The author began their journey in the startup world by co-founding Y Combinator (YC), a startup accelerator that provided funding and support to startups in batches. They aimed to fix issues in the venture capital industry by making a larger number of smaller investments, funding younger and more technical founders, and allowing founders to remain as CEOs. The author also wrote Hacker News, a news aggregator initially for startup founders, in a new version of Arc programming language.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response:  The author began their journey in the startup world by co-founding Y Combinator (YC), a startup accelerator that provided funding and support to startups in batches. They aimed to fix issues in the venture capital industry by making a larger number of smaller investments, funding younger and more technical founders, and allowing founders to remain as CEOs. The author also wrote Hacker News, a news aggregator initially for startup founders, in a new version of Arc programming language.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: Since then, the author has been involved in mentoring and advising numerous startups, helping them grow and succeed in their respective industries. [Search(What are some notable startups the author has worked with?)]\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: Since then, the author has been involved in mentoring and advising numerous startups, helping them grow and succeed in their respective industries. Some notable startups the author has worked with include Reddit, Justin Kan and Emmett Shear (who went on to found Twitch), Aaron Swartz (who helped write the RSS spec), and Sam Altman (who later became the second president of YC).\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: The author began their journey in the startup world by co-founding Y Combinator (YC), a startup accelerator that provided funding and support to startups in batches. They aimed to fix issues in the venture capital industry by making a larger number of smaller investments, funding younger and more technical founders, and allowing founders to remain as CEOs. The author also wrote Hacker News, a news aggregator initially for startup founders, in a new version of Arc programming language. Since then, the author has been involved in mentoring and advising numerous startups, helping them grow and succeed in their respective industries. Some notable startups the author has worked with include Reddit, Justin Kan and Emmett Shear (who went on to found Twitch), Aaron Swartz (who helped write the RSS spec), and Sam Altman (who later became the second president of YC).\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: done\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = flare_query_engine.query(\n",
    "    \"Can you tell me about the author's trajectory in the startup world?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "78085cc6-d9b7-486b-8c6c-5ad092678642",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author began their journey in the startup world by co-founding Y Combinator (YC), a startup accelerator that provided funding and support to startups in batches. They aimed to fix issues in the venture capital industry by making a larger number of smaller investments, funding younger and more technical founders, and allowing founders to remain as CEOs. The author also wrote Hacker News, a news aggregator initially for startup founders, in a new version of Arc programming language. Since then, the author has been involved in mentoring and advising numerous startups, helping them grow and succeed in their respective industries. Some notable startups the author has worked with include Reddit, Justin Kan and Emmett Shear (who went on to found Twitch), Aaron Swartz (who helped write the RSS spec), and Sam Altman (who later became the second president of YC). \n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6e3645ac-2220-4b32-80a6-0b82d2ec7e38",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mQuery: Can you tell me about what the author did during his time at YC?\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: During his time at YC, the author [Search(What did the author do at YC?)]\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: During his time at YC, the author worked on selecting and helping founders at YC, solving their problems, and engaging with their startups. They also wrote all of YC's internal software in Arc and managed Hacker News, which was a source of stress for them.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response:  During his time at YC, the author worked on selecting and helping founders at YC, solving their problems, and engaging with their startups. They also wrote all of YC's internal software in Arc and managed Hacker News, which was a source of stress for them.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: done\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = flare_query_engine.query(\n",
    "    \"Can you tell me about what the author did during his time at YC?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "31279b2f-7cc9-4e10-962f-e37cbe5e2d31",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "During his time at YC, the author worked on selecting and helping founders at YC, solving their problems, and engaging with their startups. They also wrote all of YC's internal software in Arc and managed Hacker News, which was a source of stress for them. \n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a5c86772-be9c-4a8f-8720-8a62ae9a7d1d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[32;1m\u001b[1;3mQuery: Tell me about the author's life from childhood to adulthood\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: \n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: The author grew up in a small town, where they [Search(What did the author do during their childhood?)] and later went on to attend college, majoring in [Search(What did the author major in during college?)].\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: The author grew up in a small town, where they mainly worked on writing and programming outside of school. They wrote short stories and tried programming on the IBM 1401 using an early version of Fortran and later went on to attend college, majoring in \n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response:  The author grew up in a small town, where they mainly worked on writing and programming outside of school. They wrote short stories and tried programming on the IBM 1401 using an early version of Fortran and later went on to attend college, majoring in\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: computer science and English literature. After college, they [Search(What did the author do after college?)]\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: computer science and English literature. After college, they wrote essays on various topics, worked on spam filters, did some painting, and hosted dinners for friends. They also bought a building in Cambridge to use as an office. Later, the author applied to art schools, got accepted into RISD, and attended their foundation classes. They also received an invitation to take the entrance exam at the Accademia di Belli Arti in Florence.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: The author grew up in a small town, where they mainly worked on writing and programming outside of school. They wrote short stories and tried programming on the IBM 1401 using an early version of Fortran and later went on to attend college, majoring in computer science and English literature. After college, they wrote essays on various topics, worked on spam filters, did some painting, and hosted dinners for friends. They also bought a building in Cambridge to use as an office. Later, the author applied to art schools, got accepted into RISD, and attended their foundation classes. They also received an invitation to take the entrance exam at the Accademia di Belli Arti in Florence.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: During their time at RISD and the Accademia di Belli Arti, the author honed their artistic skills and further developed their writing, eventually transitioning into a successful career as an author and artist. [Search(What did the author achieve in their career?)]\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mUpdated lookahead response: During their time at RISD and the Accademia di Belli Arti, the author honed their artistic skills and further developed their writing, eventually transitioning into a successful career as an author and artist. The author achieved several things in their career, including publishing essays online, writing a book called \"Hackers & Painters,\" working on spam filters, doing some painting, and hosting dinners for friends. They also discussed ideas about venture capital and how it could be improved.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mCurrent response: The author grew up in a small town, where they mainly worked on writing and programming outside of school. They wrote short stories and tried programming on the IBM 1401 using an early version of Fortran and later went on to attend college, majoring in computer science and English literature. After college, they wrote essays on various topics, worked on spam filters, did some painting, and hosted dinners for friends. They also bought a building in Cambridge to use as an office. Later, the author applied to art schools, got accepted into RISD, and attended their foundation classes. They also received an invitation to take the entrance exam at the Accademia di Belli Arti in Florence. During their time at RISD and the Accademia di Belli Arti, the author honed their artistic skills and further developed their writing, eventually transitioning into a successful career as an author and artist. The author achieved several things in their career, including publishing essays online, writing a book called \"Hackers & Painters,\" working on spam filters, doing some painting, and hosting dinners for friends. They also discussed ideas about venture capital and how it could be improved.\n",
      "\u001b[0m\u001b[38;5;200m\u001b[1;3mLookahead response: done\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = flare_query_engine.query(\n",
    "    \"Tell me about the author's life from childhood to adulthood\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b34a5e7c-37fa-4d08-b2ad-2cf757fe230a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author grew up in a small town, where they mainly worked on writing and programming outside of school. They wrote short stories and tried programming on the IBM 1401 using an early version of Fortran and later went on to attend college, majoring in computer science and English literature. After college, they wrote essays on various topics, worked on spam filters, did some painting, and hosted dinners for friends. They also bought a building in Cambridge to use as an office. Later, the author applied to art schools, got accepted into RISD, and attended their foundation classes. They also received an invitation to take the entrance exam at the Accademia di Belli Arti in Florence. During their time at RISD and the Accademia di Belli Arti, the author honed their artistic skills and further developed their writing, eventually transitioning into a successful career as an author and artist. The author achieved several things in their career, including publishing essays online, writing a book called \"Hackers & Painters,\" working on spam filters, doing some painting, and hosting dinners for friends. They also discussed ideas about venture capital and how it could be improved. \n"
     ]
    }
   ],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8f1298b6-459a-4231-814f-46a9c2629984",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = index_query_engine.query(\n",
    "    \"Can you tell me about the author's trajectory in the startup world?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b01a53e8-ec7e-4599-bcbb-532e23037083",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author's trajectory in the startup world began with their involvement in various projects and activities, such as writing essays on different topics, working on spam filters, and painting. They also hosted dinners for friends, which helped them learn how to cook for groups and network with people from various backgrounds.\n",
      "\n",
      "In October 2003, the author met Jessica Livingston at a party, who later became a significant figure in their startup journey. Jessica worked in marketing at a Boston investment bank and was intrigued by the stories of startup founders she met through the author. She decided to compile a book of interviews with these founders.\n",
      "\n",
      "In early 2005, Jessica interviewed for a marketing job at a Boston VC firm, which led the author to discuss the issues with venture capital and how it could be improved. The author also gave a talk at the Harvard Computer Society about starting a startup, which made them realize they should start angel investing.\n",
      "\n",
      "On March 11, the author, Jessica, and their friends Robert and Trevor decided to start their own investment firm, implementing the ideas they had discussed. They founded Y Combinator, an angel investment firm that made unconventional choices in the startup world. The author's trajectory in the startup world has been marked by their involvement in various projects, networking, and eventually co-founding a successful investment firm.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "a6aaa727-9007-4d0c-8386-d644c46d3a4a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "response = index_query_engine.query(\n",
    "    \"Tell me about the author's life from childhood to adulthood\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e9eb50ab-c4a8-467b-acbc-750918f721c9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author's life from childhood to adulthood includes a variety of experiences and interests. They wrote numerous essays on various topics, which were later compiled into a book called Hackers & Painters. They also worked on spam filters and pursued painting as a hobby. The author used to host dinners for friends every Thursday night, which taught them how to cook for groups. They bought a building in Cambridge, which was a former candy factory and later a porn studio, to use as an office.\n",
      "\n",
      "In October 2003, the author met Jessica Livingston at a party, and they started dating a few days later. Jessica worked in marketing at a Boston investment bank and later decided to compile a book of interviews with startup founders. When she was looking for a new job, the author shared their thoughts on how venture capital should be improved.\n",
      "\n",
      "The author also attended the Accademia, a prestigious institution, to study painting. However, they were disappointed with the lack of teaching and learning taking place there. The author painted still lives in their bedroom at night, using leftover scraps of canvas.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1277d2a0-4c1d-410f-a980-1695a970756b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
